from pathlib import Path

import pandas as pd

from ...exceptions import TdxFileNotFoundException
from ...exceptions import TdxNotAssignVipdocPathException
from ...logger import log
from ..base import BaseReader


class TdxDailyBarReader(BaseReader):
    """ 读取通达信日线数据 """

    def __init__(self, vipdoc_path=None):
        """

        :param vipdoc_path:
        """
        self.vipdoc_path = vipdoc_path

    def generate_filename(self, code, exchange):
        """

        :param code:
        :param exchange:
        :return:
        """
        if self.vipdoc_path is None:
            raise TdxNotAssignVipdocPathException(r'Please provide a vipdoc path , such as C:/newtdx/vipdoc')

        return Path(self.vipdoc_path, exchange) / 'lday' / f'{exchange}{code}.day'

    def get_kline_by_code(self, code, exchange):
        """

        :param code:
        :param exchange:
        :return:
        """
        fname = self.generate_filename(code, exchange)
        return self.parse_data_by_file(fname)

    def parse_data_by_file(self, fname):
        """

        :param fname:
        :return:
        """
        if not Path(fname).is_file():
            raise TdxFileNotFoundException('no tdx kline data, pleaes check path %s', fname)

        with open(fname, 'rb') as f:
            return self.unpack_records('<IIIIIfII', f.read())

    def get_df(self, code_or_file, exchange=None):
        """

        :param code_or_file:
        :param exchange:
        :return:
        """
        if exchange is None:
            return self.get_df_by_file(code_or_file)

        return self.get_df_by_code(code_or_file, exchange)

    def get_df_by_file(self, fname):
        """

        :param fname:
        :return:
        """
        if not Path(fname).is_file():
            raise TdxFileNotFoundException('no tdx kline data, pleaes check path %s', fname)

        security_type = self.get_security_type(fname)

        if security_type not in self.SECURITY_TYPE:
            log.debug('Unknown security type !\n')
            raise NotImplementedError

        coefficient = self.SECURITY_COEFFICIENT[security_type]
        data = [self._df_convert(row, coefficient) for row in self.parse_data_by_file(fname)]

        df = pd.DataFrame(data=data, columns=('date', 'open', 'high', 'low', 'close', 'amount', 'volume'))
        df.index = pd.to_datetime(df.date)

        return df[['open', 'high', 'low', 'close', 'amount', 'volume']]

    def get_df_by_code(self, code, exchange):
        """

        :param code:
        :param exchange:
        :return:
        """
        fname = self.generate_filename(code, exchange)
        return self.get_df_by_file(fname)

    @staticmethod
    def _df_convert(rows: list, coefficient):
        """

        :param rows:
        :param coefficient:
        :return:
        """
        t_date = str(rows[0])
        date_str = t_date[:4] + '-' + t_date[4:6] + '-' + t_date[6:]

        new_row = (
            date_str,
            rows[1] * coefficient[0],  # * 0.01 * 1000 , zipline need 1000 times to original price
            rows[2] * coefficient[0],
            rows[3] * coefficient[0],
            rows[4] * coefficient[0],
            rows[5],
            rows[6] * coefficient[1]
        )

        return new_row

    def get_security_type(self, fname):
        """

        :param fname:
        :return:
        """
        exchange = str(fname[-12:-10]).lower()
        code_head = fname[-10:-8]

        if exchange == self.SECURITY_EXCHANGE[0]:
            if code_head in ['00', '30']:
                return 'SZ_A_STOCK'

            if code_head in ['20']:
                return 'SZ_B_STOCK'

            if code_head in ['39']:
                return 'SZ_INDEX'

            if code_head in ['15', '16']:
                return 'SZ_FUND'

            if code_head in ['10', '11', '12', '13', '14']:
                return 'SZ_BOND'
        elif exchange == self.SECURITY_EXCHANGE[1]:
            if code_head in ['60', '68']:  # 688XXX科创板
                return 'SH_A_STOCK'

            if code_head in ['90']:
                return 'SH_B_STOCK'

            if code_head in ['00', '88', '99']:
                return 'SH_INDEX'

            if code_head in ['50', '51']:
                return 'SH_FUND'

            if code_head in ['01', '10', '11', '12', '13', '14', '20']:
                return 'SH_BOND'
        else:
            print('Unknown security exchange !\n')
            raise NotImplementedError

    SECURITY_EXCHANGE = ['sz', 'sh']
    SECURITY_TYPE = ['SH_A_STOCK', 'SH_B_STOCK', 'SH_INDEX', 'SH_FUND', 'SH_BOND', 'SZ_A_STOCK', 'SZ_B_STOCK', 'SZ_INDEX', 'SZ_FUND', 'SZ_BOND']
    SECURITY_COEFFICIENT = {'SH_A_STOCK': [0.01, 0.01], 'SH_B_STOCK': [0.001, 0.01], 'SH_INDEX': [0.01, 1.0], 'SH_FUND': [0.001, 1.0], 'SH_BOND': [0.001, 1.0], 'SZ_A_STOCK': [0.01, 0.01],
                            'SZ_B_STOCK': [0.01, 0.01], 'SZ_INDEX': [0.01, 1.0], 'SZ_FUND': [0.001, 0.01], 'SZ_BOND': [0.001, 0.01]}
